package com.iussoft.gateway.filter;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.UnpooledByteBufAllocator;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @description: 缓存请求参数，解决body只能读一次问题
 * @author: Hou Dayu
 * @date: Created in 2020/2/29
 */
public class RecorderServerHttpRequestDecorator extends ServerHttpRequestDecorator {

    private boolean bufferCached = false;
    private Mono<Void> progress = null;
    private final List<DataBuffer> dataBuffers = new ArrayList<>();

    public RecorderServerHttpRequestDecorator(ServerHttpRequest delegate) {
        super(delegate);
    }

    @Override
   public Flux<DataBuffer> getBody() {
       synchronized (dataBuffers) {
           if (bufferCached)
               return copy();

           if (progress == null) {
               progress = cache();
           }
           return progress.thenMany(Flux.defer(this::copy));
       }
   }

    private Flux<DataBuffer> copy() {
        return Flux.fromIterable(dataBuffers).map(buf -> buf.factory().wrap(buf.asByteBuffer()));
    }

    private Mono<Void> cache() {
        return super.getBody()
                .map(dataBuffers::add)
                .then(Mono.defer(()-> {
                    bufferCached = true;
                    progress = null;
                    return Mono.empty();
                }));
    }

    /**
     * 请求body参数内容
     * @return
     */
    public AtomicReference<String> getRequestBody(){
        Flux<DataBuffer> body = getBody();
        AtomicReference<String> requestBody = new AtomicReference<>("");
        body.subscribe(buffer -> { //读取requestBody传参
            CharBuffer charBuffer = StandardCharsets.UTF_8.decode(buffer.asByteBuffer());
            requestBody.set(charBuffer.toString());
            DataBufferUtils.release(buffer);
        });
        return requestBody;
    }
}
